package com.baidb.gateway.ribbon.rule;

import com.netflix.loadbalancer.ILoadBalancer;
import com.netflix.loadbalancer.RoundRobinRule;
import com.netflix.loadbalancer.Server;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.atomic.AtomicInteger;
import org.springframework.cloud.consul.discovery.ConsulServer;

/**
 * @author baidb
 * @date 2021/6/17.
 */
public class RibbonServerDiscoveryRule extends RoundRobinRule {

    private final AtomicInteger nextServerCyclicCounter = new AtomicInteger(0);

    public RibbonServerDiscoveryRule() {
        super();
    }

    public RibbonServerDiscoveryRule(ILoadBalancer lb) {
        super(lb);
    }

    @Override
    public Server choose(ILoadBalancer lb, Object key) {
        if (lb == null) {
            return null;
        }
        Server server = null;
        int count = 0;
        while (server == null && count++ < 10) {
            List<Server> reachableServers = lb.getReachableServers();
            List<Server> allServers = lb.getAllServers();
            int upCount = reachableServers.size();
            int serverCount = allServers.size();
            if ((upCount == 0) || (serverCount == 0)) {
                return null;
            }
            int nextServerIndex = incrementAndGetModulo(serverCount);
            server = allServers.get(nextServerIndex);
            if (server == null) {
                Thread.yield();
                continue;
            }
            if (server.isAlive() && (server.isReadyToServe()) && available(server, key)) {
                return (server);
            }
            server = null;
        }
        return server;
    }

    @SuppressWarnings("unchecked")
    private boolean available(Server server, Object key) {
        if (key instanceof Map) {
            String version = ((Map<String, String>)key).get("version");
            if (Objects.nonNull(version)) {
                return Objects.equals(((ConsulServer)server).getMetadata().get("version"), version);
            }
        }
        return true;
    }

    private int incrementAndGetModulo(int modulo) {
        for (; ; ) {
            int current = nextServerCyclicCounter.get();
            int next = (current + 1) % modulo;
            if (nextServerCyclicCounter.compareAndSet(current, next)) {
                return next;
            }
        }
    }

}
